import uuid
import os
import re
import json
import tempfile
import time
import traceback
import sys
import asyncio
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.responses import JSONResponse
from typing import List, Optional
from paddleocr import PaddleOCR
from PIL import Image
import paddle
import paddle.distributed as dist
import numpy as np

# # 方法1: 使用数据并行
# strategy = dist.Strategy()
# # 设置要使用的设备列表
# strategy.npu_config.threads_per_core = 1
# strategy.npu_config.npus = [0, 1, 2, 3]  # 指定要使用的NPU设备ID
#
# # 初始化并行环境
# dist.init_parallel_env()
# Import PyMuPDF (fitz module)
# Note: There's a conflicting 'fitz' package that may cause issues.
# If you get import errors, uninstall the 'fitz' package and ensure PyMuPDF is installed:
# pip uninstall fitz
# pip install PyMuPDF
try:
    import fitz  # PyMuPDF
except ImportError as e:
    raise ImportError(
        "Failed to import PyMuPDF (fitz module). "
        "Please ensure PyMuPDF is installed: pip install PyMuPDF\n"
        "If you have a conflicting 'fitz' package, uninstall it first: pip uninstall fitz\n"
        f"Original error: {e}"
    )
except Exception as e:
    # Handle RuntimeError from conflicting fitz package
    if "Directory 'static/' does not exist" in str(e):
        raise ImportError(
            "Conflicting 'fitz' package detected. This is not PyMuPDF.\n"
            "Please uninstall the conflicting package and install PyMuPDF:\n"
            "  pip uninstall fitz\n"
            "  pip install PyMuPDF\n"
            f"Original error: {e}"
        )
    raise

# 初始化 FastAPI 应用
app = FastAPI(title="OCR识别API (PPStructureV3)", description="支持文件上传进行OCR识别，使用PPStructureV3版本")

# 全局初始化 OCR 模型（避免重复加载）
ocr_instance = None

# 全局信号量：限制并发OCR任务数量（避免模型并发访问冲突）
# PaddleOCR 模型实例不是线程安全的，必须使用信号量保护
# 设置为None，在首次使用时延迟初始化（确保在正确的事件循环中创建）
ocr_semaphore = None


def init_semaphore():
    """初始化信号量（延迟初始化，确保在事件循环中创建）"""
    global ocr_semaphore
    if ocr_semaphore is None:
        ocr_semaphore = asyncio.Semaphore(1)
    return ocr_semaphore


def init_ocr_model():
    """初始化 OCR 模型"""
    global ocr_instance
    if ocr_instance is None:
        try:
            # paddle.set_device('npu:1')
            #print("✓ 已设置设备为 npu:0")
            ocr_instance = PaddleOCR(
                ocr_version="PP-OCRv3",
                # device="npu:2,3",  # 使用2号和3号NPU设备
                #use_doc_orientation_classify=False,  # 禁用文档方向分类模块
                #use_doc_unwarping=False,  # 禁用文本图像矫正模块
                # use_table_recognition=True,    # 启用表格识别
                # use_region_detection=True,     # 启用区域检测
                #lang="ch"                      # 中文模型
            )
        except TypeError as e:
            # 参数类型错误（比如参数名错误、参数类型不匹配等）
            error_msg = str(e)
            error_trace = traceback.format_exc()
            exception_type = type(e).__name__
            
            detailed_error_parts = [
                f"OCR模型初始化失败 - 参数错误:",
                f"",
                f"异常类型: {exception_type}",
                f"错误消息: {error_msg}",
                f"",
                f"当前使用的参数:",
                f"  ocr_version = 'PP-OCRv3'",
                f"  use_doc_orientation_classify = False",
                f"  use_doc_unwarping = False",
                f"",
                f"可能的问题:",
                f"  - PaddleOCR 初始化失败",
                f"  - 请检查 PaddleOCR 的文档，确认正确的参数名称",
                f"  - 检查依赖是否正确安装",
                f"",
                f"完整堆栈跟踪:",
                f"{error_trace}"
            ]
            
            detailed_error = "\n".join(detailed_error_parts)
            
            # 强制刷新输出，确保错误信息立即打印
            sys.stdout.flush()
            sys.stderr.flush()
            
            print("=" * 80, flush=True)
            print("OCR模型初始化失败 - 参数错误:", flush=True)
            print("=" * 80, flush=True)
            print(detailed_error, flush=True)
            print("=" * 80, flush=True)
            
            raise HTTPException(
                status_code=500,
                detail=detailed_error
            )
        except Exception as e:
            error_msg = str(e)
            error_trace = traceback.format_exc()
            
            # 提取异常类型和参数
            exception_type = type(e).__name__
            exception_args = str(e.args) if hasattr(e, 'args') and e.args else ""
            
            # 强制刷新输出
            sys.stdout.flush()
            sys.stderr.flush()
            
            # 检查是否是 paddlex 的 DependencyError（需要额外依赖）
            if "DependencyError" in exception_type or ("paddlex" in error_msg.lower() and "additional dependencies" in error_msg.lower()):
                # 提取 paddlex 版本（如果可以从堆栈中提取）
                paddlex_version = "3.3.6"  # 默认版本，可以从 requirements 文件中读取
                if "PADDLEX_VERSION" in error_msg:
                    # 尝试从错误消息中提取版本号
                    version_match = re.search(r'==([\d.]+)', error_trace)
                    if version_match:
                        paddlex_version = version_match.group(1)
                
                detailed_error_parts = [
                    f"PP-StructureV3 依赖错误 (paddlex 需要 OCR 额外依赖):",
                    f"",
                    f"异常类型: {exception_type}",
                    f"错误消息: {error_msg}",
                    f"",
                    f"解决方案:",
                    f"",
                    f"1. 卸载当前的 paddlex:",
                    f"   pip uninstall paddlex",
                    f"",
                    f"2. 安装带 OCR 额外依赖的 paddlex:",
                    f"   pip install \"paddlex[ocr]=={paddlex_version}\"",
                    f"",
                    f"或者直接运行（会自动安装正确版本）:",
                    f"   pip install -r requirements_linux.txt",
                    f"",
                    f"注意: requirements_linux.txt 已更新为 paddlex[ocr]=={paddlex_version}",
                    f"",
                    f"完整堆栈跟踪:",
                    f"{error_trace}"
                ]
                
                detailed_error = "\n".join(detailed_error_parts)
                
                # 打印到控制台（便于日志查看）
                print("=" * 80, flush=True)
                print("OCR模型初始化失败 - paddlex OCR 依赖缺失:", flush=True)
                print("=" * 80, flush=True)
                print(detailed_error, flush=True)
                print("=" * 80, flush=True)
                
                # 将完整错误信息包含在HTTP响应中
                raise HTTPException(
                    status_code=500,
                    detail=detailed_error
                )
            
            # 检查是否是依赖错误
            elif "dependency" in error_msg.lower() or "dependencies" in error_msg.lower():
                # 尝试提取更详细的错误信息
                detailed_error_parts = [
                    f"依赖错误详情:",
                    f"",
                    f"异常类型: {exception_type}",
                    f"错误消息: {error_msg}",
                    f"",
                    f"完整堆栈跟踪:",
                    f"{error_trace}",
                    f"",
                    f"请检查以下依赖是否已正确安装:",
                    f"  - paddleocr (pip install paddleocr)",
                    f"  - paddlepaddle (pip install paddlepaddle)",
                    f"  - paddlex[ocr] (pip install \"paddlex[ocr]==3.3.6\")",
                    f"  - 如果使用NPU设备: paddlepaddle-npu",
                    f"  - 如果使用GPU设备: paddlepaddle-gpu",
                    f"  - 其他相关依赖（请参考 paddleocr 官方文档）",
                    f"",
                    f"提示: 请查看上面的堆栈跟踪，找到具体是哪个模块导入失败"
                ]
                
                detailed_error = "\n".join(detailed_error_parts)
                
                # 打印到控制台（便于日志查看）
                print("=" * 80, flush=True)
                print("OCR模型初始化失败 - 依赖错误:", flush=True)
                print("=" * 80, flush=True)
                print(detailed_error, flush=True)
                print("=" * 80, flush=True)
                
                # 将完整错误信息包含在HTTP响应中
                raise HTTPException(
                    status_code=500,
                    detail=detailed_error
                )
            else:
                # 其他类型的错误
                detailed_error_parts = [
                    f"OCR模型初始化失败:",
                    f"",
                    f"异常类型: {exception_type}",
                    f"错误消息: {error_msg}",
                    f"",
                    f"异常参数: {exception_args}" if exception_args else "",
                    f"",
                    f"完整堆栈跟踪:",
                    f"{error_trace}"
                ]
                
                detailed_error = "\n".join([part for part in detailed_error_parts if part])
                
                # 打印到控制台
                print("=" * 80, flush=True)
                print("OCR模型初始化失败:", flush=True)
                print("=" * 80, flush=True)
                print(detailed_error, flush=True)
                print("=" * 80, flush=True)
                
                # 同时打印到 stderr
                print(detailed_error, file=sys.stderr, flush=True)
                
                # 将完整错误信息包含在HTTP响应中
                raise HTTPException(
                    status_code=500,
                    detail=detailed_error
                )
    return ocr_instance


def pdf_to_images(pdf_path, output_dir, dpi=200):
    """
    将PDF转换为图片列表

    Args:
        pdf_path: PDF文件路径
        output_dir: 输出目录
        dpi: 转换分辨率（默认200）

    Returns:
        tuple: (图片路径列表, 转换耗时(秒))
    """
    image_paths = []
    start_time = time.time()
    try:
        pdf_document = fitz.open(pdf_path)
        total_pages = len(pdf_document)
        
        for page_num in range(total_pages):
            page = pdf_document[page_num]
            # 设置缩放比例（DPI）
            zoom = dpi / 72.0
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)

            # 保存为PNG图片
            image_path = output_dir / f"page_{page_num + 1}.png"
            pix.save(image_path)
            image_paths.append(image_path)

        pdf_document.close()
        
        elapsed_time = time.time() - start_time
        print("=" * 80)
        print(f"PDF转换完成:")
        print(f"  总页数: {total_pages}")
        print(f"  转换耗时: {elapsed_time:.2f} 秒 ({elapsed_time/60:.2f} 分钟)")
        print(f"  平均每页: {elapsed_time/total_pages:.2f} 秒" if total_pages > 0 else "")
        print("=" * 80)
        
        return image_paths, elapsed_time
    except Exception as e:
        elapsed_time = time.time() - start_time
        print(f"PDF转换失败 (耗时 {elapsed_time:.2f} 秒): {e}")
        print(traceback.format_exc())
        raise HTTPException(status_code=500, detail=f"PDF转换失败: {str(e)}")


def convert_numpy_to_list(obj):
    """
    递归地将对象中的所有 numpy 数组和非JSON可序列化对象转换为可序列化格式
    
    Args:
        obj: 任意对象（字典、列表、numpy数组等）
    
    Returns:
        转换后的对象，所有不可序列化对象都被转换为可序列化格式
    """
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {key: convert_numpy_to_list(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [convert_numpy_to_list(item) for item in obj]
    elif isinstance(obj, tuple):
        return tuple(convert_numpy_to_list(item) for item in obj)
    elif isinstance(obj, (str, int, float, bool, type(None))):
        # 基本类型直接返回
        return obj
    else:
        # 对于其他不可序列化的对象（如Font对象等），转换为字符串
        try:
            import json
            json.dumps(obj)  # 尝试序列化
            return obj
        except (TypeError, ValueError):
            # 如果无法序列化，转换为字符串表示
            return str(obj)


def preprocess_image(image_path, max_size=(4000, 4000), max_megapixels=16, enable_compression=False):
    """
    预处理图片：根据配置决定是否压缩图片

    Args:
        image_path: 图片路径
        max_size: 最大尺寸 (width, height)
        max_megapixels: 最大百万像素数（16MP = 约4000x4000）
        enable_compression: 压缩控制参数
            - True: 强制压缩所有图片
            - False: 只有图片过大时才自动压缩（默认）

    Returns:
        tuple: (处理后的图片路径, 是否压缩, 原始文件大小(字节), 压缩后文件大小(字节))
    """
    try:
        original_size = image_path.stat().st_size if image_path.exists() else 0
        
        with Image.open(image_path) as img:
            width, height = img.size
            megapixels = (width * height) / 1_000_000

            # 判断是否需要压缩
            need_compress = False
            if enable_compression:
                # 强制压缩模式：所有图片都压缩
                need_compress = True
            else:
                # 自动压缩模式：只有图片过大时才压缩
                if megapixels > max_megapixels or width > max_size[0] or height > max_size[1]:
                    need_compress = True
            
            if need_compress:
                # 计算缩放比例
                scale = min(max_size[0] / width, max_size[1] / height,
                            (max_megapixels * 1_000_000 / (width * height)) ** 0.5)
                new_width = int(width * scale)
                new_height = int(height * scale)

                # 压缩图片
                img_resized = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

                # 保存压缩后的图片
                processed_path = image_path.parent / f"{image_path.stem}_processed{image_path.suffix}"
                img_resized.save(processed_path, quality=95, optimize=True)
                
                compressed_size = processed_path.stat().st_size if processed_path.exists() else 0

                return processed_path, True, original_size, compressed_size  # 返回新路径、是否压缩标志、原始大小、压缩后大小
            else:
                return image_path, False, original_size, original_size  # 原路径，未压缩，大小相同
    except Exception as e:
        print(f"图片预处理失败，使用原图: {e}")
        original_size = image_path.stat().st_size if image_path.exists() else 0
        return image_path, False, original_size, original_size


def process_single_file_ocr(file_content: bytes, filename: str, file_id: str = None, enable_compression: bool = False):
    """
    处理单个文件的OCR识别
    
    Args:
        file_content: 文件内容（字节）
        filename: 文件名
        file_id: 文件ID（可选，如果不提供则自动生成）
        enable_compression: 是否强制压缩图片（True=强制压缩所有图片，False=只在图片过大时自动压缩）
    
    Returns:
        dict: OCR识别结果
    """
    # 使用当前目录下的 tmp 文件夹
    base_dir = Path(__file__).parent.parent  # 获取项目根目录
    tmp_dir = base_dir / "tmp"
    tmp_dir.mkdir(exist_ok=True)

    # 生成唯一文件名
    if file_id is None:
        file_id = str(uuid.uuid4())
    file_ext = Path(filename).suffix.lower() if filename else ''
    input_file_path = tmp_dir / f"{file_id}{file_ext}"
    processed_file_paths = []  # 存储所有处理后的文件路径（包括PDF转换的图片）
    was_compressed = False
    original_image_size = 0  # 原始图片大小（字节）
    compressed_image_size = 0  # 压缩后图片大小（字节）
    is_pdf = False
    pdf_pages = 0

    try:
        # 保存文件内容
        with open(input_file_path, "wb") as buffer:
            buffer.write(file_content)

        # 初始化 OCR 模型
        ocr = init_ocr_model()

        # 判断文件类型并处理
        if file_ext == '.pdf':
            # PDF文件：转换为图片后逐页OCR
            is_pdf = True
            pdf_images_dir = tmp_dir / f"{file_id}_pages"
            pdf_images_dir.mkdir(exist_ok=True)

            # 将PDF转换为图片（记录时间）
            image_paths, pdf_conversion_time = pdf_to_images(input_file_path, pdf_images_dir)
            pdf_pages = len(image_paths)

            # 对所有页面进行OCR识别
            all_results = []
            all_texts = []
            page_results = []
            total_ocr_time = 0.0  # 总OCR耗时
            page_ocr_times = []  # 每页OCR耗时

            for page_idx, image_path in enumerate(image_paths):
                try:
                    # 预处理图片（压缩）
                    processed_path, compressed, orig_size, comp_size = preprocess_image(image_path, enable_compression=enable_compression)
                    if compressed:
                        processed_file_paths.append(processed_path)
                        # 记录压缩信息（使用第一个压缩的图片的信息，或累计）
                        if not was_compressed:
                            original_image_size = orig_size
                            compressed_image_size = comp_size
                            was_compressed = True
                    # 始终记录原始图片路径以便删除
                    processed_file_paths.append(image_path)

                    # 执行OCR识别（记录时间）
                    ocr_input_path = str(processed_path) if compressed else str(image_path)
                    page_ocr_start = time.time()
                    try:
                        output = ocr.predict(ocr_input_path)
                    except Exception as ocr_error:
                        error_msg = str(ocr_error)
                        error_trace = traceback.format_exc()
                        exception_type = type(ocr_error).__name__
                        
                        # 检查是否是依赖错误
                        if "dependency" in error_msg.lower() or "dependencies" in error_msg.lower():
                            detailed_error = (
                                f"第 {page_idx + 1} 页OCR识别失败 - 依赖错误:\n"
                                f"\n异常类型: {exception_type}\n"
                                f"错误消息: {error_msg}\n"
                                f"\n完整堆栈跟踪:\n{error_trace}\n"
                                f"\n请检查依赖是否已正确安装"
                            )
                            print("=" * 80)
                            print(detailed_error)
                            print("=" * 80)
                            raise Exception(detailed_error)
                        else:
                            raise
                    
                    page_ocr_time = time.time() - page_ocr_start
                    total_ocr_time += page_ocr_time
                    page_ocr_times.append({
                        "page": page_idx + 1,
                        "time": page_ocr_time
                    })
                    
                    print(f"第 {page_idx + 1} 页OCR识别完成，耗时: {page_ocr_time:.2f} 秒")

                    # 处理PaddleOCR的结果
                    page_texts = []
                    for res in output:
                        # PaddleOCR返回的是字典，包含rec_texts字段
                        if isinstance(res, dict):
                            rec_texts = res.get("rec_texts", [])
                            if rec_texts:
                                page_texts.extend(rec_texts)
                        elif hasattr(res, 'rec_texts'):
                            # 如果是对象，尝试获取rec_texts属性
                            rec_texts = res.rec_texts if isinstance(res.rec_texts, list) else [res.rec_texts]
                            page_texts.extend([str(t) for t in rec_texts if t])
                        else:
                            # 尝试其他可能的文本字段
                            if hasattr(res, 'text'):
                                page_texts.append(str(res.text))
                            elif isinstance(res, str):
                                page_texts.append(res)

                        # 保存结果
                        all_results.append(res)

                    # 提取所有文本
                    all_texts.extend(page_texts)
                    page_full_text = "\n".join(page_texts)

                    # 记录每页的结果
                    page_results.append({
                        "page": page_idx + 1,
                        "texts": page_texts,
                        "full_text": page_full_text,
                        "text_count": len(page_texts)
                    })
                except Exception as e:
                    print(f"第 {page_idx + 1} 页OCR失败: {e}")
                    page_results.append({
                        "page": page_idx + 1,
                        "error": str(e),
                        "texts": [],
                        "full_text": "",
                        "text_count": 0
                    })

            # 格式化最终结果（PaddleOCR格式）
            ocr_results = []
            for idx, res in enumerate(all_results):
                # PaddleOCR返回的是字典格式
                if isinstance(res, dict):
                    rec_texts = res.get("rec_texts", [])
                    result_item = {
                        "result_index": idx,
                        "rec_texts": rec_texts,
                        "full_text": "\n".join(rec_texts) if rec_texts else ""
                    }
                    # 只添加必要的字段
                    if "dt_boxes" in res:
                        result_item["dt_boxes"] = res["dt_boxes"]
                    if "rec_scores" in res:
                        result_item["rec_scores"] = res["rec_scores"]
                else:
                    result_item = {
                        "result_index": idx,
                        "raw_result": str(res)
                    }
                ocr_results.append(result_item)

            full_text = "\n".join(all_texts)
            
            # 打印PDF OCR总耗时
            print("=" * 80)
            print(f"PDF OCR识别完成:")
            print(f"  总页数: {pdf_pages}")
            print(f"  PDF转换耗时: {pdf_conversion_time:.2f} 秒 ({pdf_conversion_time/60:.2f} 分钟)")
            print(f"  OCR识别总耗时: {total_ocr_time:.2f} 秒 ({total_ocr_time/60:.2f} 分钟)")
            print(f"  平均每页OCR: {total_ocr_time/pdf_pages:.2f} 秒" if pdf_pages > 0 else "")
            print(f"  总处理时间: {pdf_conversion_time + total_ocr_time:.2f} 秒 ({(pdf_conversion_time + total_ocr_time)/60:.2f} 分钟)")
            print("=" * 80)

        else:
            # 图片文件处理
            # 如果是图片文件，进行预处理（压缩大图片）
            if file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp']:
                processed_file_path, was_compressed, original_image_size, compressed_image_size = preprocess_image(input_file_path, enable_compression=enable_compression)
                if was_compressed:
                    processed_file_paths.append(processed_file_path)
                ocr_input_path = str(processed_file_path) if was_compressed else str(input_file_path)
            else:
                ocr_input_path = str(input_file_path)

            # 执行 OCR 识别（记录时间）
            ocr_start_time = time.time()
            try:
                output = ocr.predict(ocr_input_path)
                
                ocr_time = time.time() - ocr_start_time
                print("=" * 80)
                print(f"图片OCR识别完成:")
                print(f"  识别耗时: {ocr_time:.2f} 秒 ({ocr_time/60:.2f} 分钟)")
                print("=" * 80)
                
            except Exception as ocr_error:
                ocr_time = time.time() - ocr_start_time
                error_msg = str(ocr_error)
                error_trace = traceback.format_exc()
                exception_type = type(ocr_error).__name__
                
                # 检查是否是依赖错误
                if "dependency" in error_msg.lower() or "dependencies" in error_msg.lower():
                    detailed_error = (
                        f"OCR识别失败 - 依赖错误 (耗时 {ocr_time:.2f} 秒):\n"
                        f"\n异常类型: {exception_type}\n"
                        f"错误消息: {error_msg}\n"
                        f"\n完整堆栈跟踪:\n{error_trace}\n"
                        f"\n请检查依赖是否已正确安装"
                    )
                    print("=" * 80)
                    print(detailed_error)
                    print("=" * 80)
                    raise Exception(detailed_error)
                
                # 如果是内存相关错误，给出更友好的提示
                if "memory" in error_msg.lower() or "内存" in error_msg or "OOM" in error_msg:
                    print(f"内存错误 (耗时 {ocr_time:.2f} 秒): {error_msg}")
                    raise Exception("内存不足，图片太大或文本过长。已自动压缩图片，如仍有问题请尝试更小的图片。")
                
                # 其他错误
                detailed_error = (
                    f"OCR识别失败 (耗时 {ocr_time:.2f} 秒):\n"
                    f"\n异常类型: {exception_type}\n"
                    f"错误消息: {error_msg}\n"
                    f"\n完整堆栈跟踪:\n{error_trace}"
                )
                print("=" * 80)
                print(detailed_error)
                print("=" * 80)
                raise Exception(detailed_error)

            # 格式化OCR结果（PaddleOCR格式）
            ocr_results = []
            all_texts = []

            for res in output:
                # PaddleOCR返回的是字典，包含rec_texts字段
                if isinstance(res, dict):
                    rec_texts = res.get("rec_texts", [])
                    all_texts.extend(rec_texts)
                    
                    result_item = {
                        "rec_texts": rec_texts,  # 识别文本列表
                        "full_text": "\n".join(rec_texts) if rec_texts else ""  # 完整文本（换行连接）
                    }
                    
                    # 只添加必要的字段（坐标、置信度等）
                    if "dt_boxes" in res:
                        result_item["dt_boxes"] = res["dt_boxes"]
                    if "rec_scores" in res:
                        result_item["rec_scores"] = res["rec_scores"]
                else:
                    # 如果不是字典，尝试获取所有属性
                    result_item = {"raw_result": str(res)}
                    # 尝试提取文本
                    if hasattr(res, 'rec_texts'):
                        rec_texts = res.rec_texts if isinstance(res.rec_texts, list) else [res.rec_texts]
                        all_texts.extend([str(t) for t in rec_texts if t])
                        result_item["rec_texts"] = [str(t) for t in rec_texts if t]
                    elif hasattr(res, 'text'):
                        text = str(res.text)
                        all_texts.append(text)
                        result_item["rec_texts"] = [text]

                ocr_results.append(result_item)

            full_text = "\n".join(all_texts)
            page_results = None  # 图片文件没有分页结果

        # 检查结果大小，如果文本过长可能存在问题
        if len(full_text) > 100000:  # 如果文本超过10万字符，给出警告
            print(f"警告: 识别文本过长 ({len(full_text)} 字符)")

        # 打印识别的结果
        print("=" * 80)
        print("OCR识别结果 (full_text):")
        print("=" * 80)
        print(full_text)
        print("=" * 80)
        
        # 如果压缩了图片，打印压缩前后的大小
        if was_compressed:
            def format_size(size_bytes):
                """格式化文件大小"""
                for unit in ['B', 'KB', 'MB', 'GB']:
                    if size_bytes < 1024.0:
                        return f"{size_bytes:.2f} {unit}"
                    size_bytes /= 1024.0
                return f"{size_bytes:.2f} TB"
            
            print(f"图片已压缩:")
            print(f"  压缩前大小: {format_size(original_image_size)} ({original_image_size} 字节)")
            print(f"  压缩后大小: {format_size(compressed_image_size)} ({compressed_image_size} 字节)")
            compression_ratio = (1 - compressed_image_size / original_image_size) * 100 if original_image_size > 0 else 0
            print(f"  压缩率: {compression_ratio:.2f}%")

        # 返回结果
        response_data = {
            "status": "success",
            "file_id": file_id,
            "original_filename": filename,
            "file_type": "pdf" if is_pdf else "image",
            "results": ocr_results,
            "all_texts": all_texts,  # 所有识别文本的列表
            "full_text": full_text,  # 完整文本（所有文本用换行连接）
            "text_count": len(all_texts),
            "result_count": len(ocr_results),
            "image_compressed": was_compressed,  # 是否压缩了图片
            "text_length": len(full_text)  # 文本总长度
        }

        # 如果压缩了图片，添加压缩信息到响应中
        if was_compressed:
            response_data["original_image_size"] = original_image_size  # 原始大小（字节）
            response_data["compressed_image_size"] = compressed_image_size  # 压缩后大小（字节）
            response_data["compression_ratio"] = round((1 - compressed_image_size / original_image_size) * 100, 2) if original_image_size > 0 else 0  # 压缩率

        # PDF特有信息
        if is_pdf:
            response_data["pdf_pages"] = pdf_pages
            response_data["page_results"] = page_results

        # 转换所有 numpy 数组为列表，确保可以 JSON 序列化
        response_data = convert_numpy_to_list(response_data)
        
        return response_data

    finally:
        # 删除上传的文件和处理后的文件（包括PDF转换的图片）
        try:
            if input_file_path.exists():
                input_file_path.unlink()

            # 删除所有处理后的文件
            for file_path in processed_file_paths:
                if file_path and file_path.exists():
                    file_path.unlink()

            # 删除PDF页面图片目录
            if is_pdf:
                pdf_images_dir = tmp_dir / f"{file_id}_pages"
                if pdf_images_dir.exists():
                    import shutil
                    shutil.rmtree(pdf_images_dir, ignore_errors=True)
        except Exception as e:
            print(f"删除文件失败: {e}")


@app.get("/")
async def root():
    """根路径，返回 API 信息"""
    return {"message": "OCR识别API (PPStructureV3)", "status": "running", "version": "PPStructureV3"}


@app.post("/ocr")
async def upload_and_ocr(
    file: UploadFile = File(...),
    enable_compression: bool = Form(False)
):
    """
    上传文件并进行OCR识别（支持图片和PDF扫描件）

    Args:
        file: 上传的文件（支持图片和PDF）
        enable_compression: 是否强制压缩图片（True=强制压缩所有图片，False=只在图片过大时自动压缩）

    Returns:
        JSONResponse: 包含OCR识别结果的响应
    """
    try:
        # 读取文件内容
        content = await file.read()
        
        # 调用核心处理函数
        result = process_single_file_ocr(content, file.filename, enable_compression=enable_compression)
        
        # 返回full_text字段（保持原有行为）
        return JSONResponse(result["full_text"])
    
    except HTTPException:
        # 重新抛出HTTP异常
        raise
    except Exception as e:
        error_msg = str(e)
        error_trace = traceback.format_exc()
        exception_type = type(e).__name__
        
        # 检查是否是依赖错误
        if "dependency" in error_msg.lower() or "dependencies" in error_msg.lower():
            detailed_error = (
                f"OCR处理失败 - 依赖错误:\n"
                f"\n异常类型: {exception_type}\n"
                f"错误消息: {error_msg}\n"
                f"\n完整堆栈跟踪:\n{error_trace}\n"
                f"\n请检查依赖是否已正确安装"
            )
            print("=" * 80)
            print(detailed_error)
            print("=" * 80)
            raise HTTPException(
                status_code=500,
                detail=detailed_error
            )
        else:
            detailed_error = (
                f"OCR处理失败:\n"
                f"\n异常类型: {exception_type}\n"
                f"错误消息: {error_msg}\n"
                f"\n完整堆栈跟踪:\n{error_trace}"
            )
            print("=" * 80)
            print(detailed_error)
            print("=" * 80)
            raise HTTPException(
                status_code=500,
                detail=detailed_error
            )


async def process_single_file_async(content: bytes, filename: str, idx: int, total: int, enable_compression: bool = False):
    """
    异步处理单个文件的OCR识别（包装同步函数）
    
    Args:
        content: 文件内容（字节）
        filename: 文件名
        idx: 文件序号（用于日志）
        total: 文件总数（用于日志）
        enable_compression: 是否强制压缩图片
    
    Returns:
        str: OCR识别的文本结果或错误信息
    """
    file_start_time = time.time()
    print(f"\n[异步任务 {idx}/{total}] 排队等待: {filename}")
    
    try:
        # 初始化信号量（延迟初始化，确保在事件循环中创建）
        semaphore = init_semaphore()
        
        # 使用信号量控制并发访问OCR模型
        # 关键：PaddleOCR 模型实例不是线程安全的，多线程同时访问会导致 Tensor 状态冲突
        async with semaphore:
            print(f"[异步任务 {idx}/{total}] 获得执行权，开始处理: {filename}")
            
            # 使用 asyncio.to_thread 将同步函数转换为异步执行
            # 这样可以避免阻塞事件循环
            result = await asyncio.to_thread(
                process_single_file_ocr,
                content,
                filename,
                None,  # file_id
                enable_compression
            )
            
            file_time = time.time() - file_start_time
            print(f"[异步任务 {idx}/{total}] 文件 {filename} 处理完成，耗时: {file_time:.2f} 秒")
            
            # 只返回full_text字符串
            return result["full_text"]
    
    except Exception as e:
        error_msg = str(e)
        file_time = time.time() - file_start_time
        print(f"[异步任务 {idx}/{total}] 文件 {filename} 处理失败 (耗时 {file_time:.2f} 秒): {error_msg}")
        
        # 失败时返回错误提示
        return f"[处理失败: {filename}]"


@app.post("/ocr_batch")
async def upload_and_ocr_batch(
    files: List[UploadFile] = File(...),
    enable_compression: bool = Form(False)
):
    """
    批量上传文件并进行OCR识别（支持图片和PDF扫描件）- 支持异步并发处理
    
    本接口采用异步并发处理模式：
    - 每个任务启动间隔1秒，避免瞬间资源占用过高
    - 使用信号量保护OCR模型访问（PaddleOCR不是线程安全的）
    - 任务异步排队执行，接口保持响应不阻塞
    - 使用线程池执行OCR任务，不阻塞事件循环
    
    注意：虽然任务是异步的，但由于模型线程安全限制，同一时间只能有1个任务在执行OCR。
         多个任务会自动排队，按顺序处理。

    Args:
        files: 上传的文件列表（支持图片和PDF）
        enable_compression: 是否强制压缩图片（True=强制压缩所有图片，False=只在图片过大时自动压缩）

    Returns:
        JSONResponse: 包含所有文件OCR识别结果的列表
    """
    batch_start_time = time.time()
    
    print("=" * 80)
    print(f"批量OCR处理开始，共 {len(files)} 个文件")
    print(f"处理模式: 异步并发（启动间隔1秒）")
    print(f"压缩模式: {'强制压缩' if enable_compression else '自动压缩'}")
    print("=" * 80)
    
    # 准备所有任务
    tasks = []
    
    for idx, file in enumerate(files, 1):
        # 读取文件内容
        content = await file.read()
        
        # 创建异步任务
        task = process_single_file_async(content, file.filename, idx, len(files), enable_compression)
        tasks.append(task)
        
        # 延迟1秒再启动下一个任务（最后一个文件不需要延迟）
        if idx < len(files):
            await asyncio.sleep(1)
            print(f"[调度器] 已启动任务 {idx}，等待1秒后启动下一个...")
    
    print(f"\n[调度器] 所有 {len(files)} 个任务已启动，等待完成...")
    
    # 并发执行所有任务，等待所有结果
    results = await asyncio.gather(*tasks)
    
    batch_time = time.time() - batch_start_time
    
    print("=" * 80)
    print(f"批量OCR处理完成:")
    print(f"  文件总数: {len(files)}")
    print(f"  成功: {sum(1 for r in results if not r.startswith('[处理失败'))}")
    print(f"  失败: {sum(1 for r in results if r.startswith('[处理失败'))}")
    print(f"  总耗时: {batch_time:.2f} 秒 ({batch_time/60:.2f} 分钟)")
    print(f"  平均每个文件: {batch_time/len(files):.2f} 秒" if len(files) > 0 else "")
    print("=" * 80)
    
    return JSONResponse(results)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8869)

